from langchain_community.vectorstores import Qdrant
from langchain_core.output_parsers import StrOutputParser
from langchain_core.prompts import ChatPromptTemplate
from langchain_core.runnables import RunnableParallel, RunnablePassthrough
from langchain_community.embeddings import ZhipuAIEmbeddings
from langchain_community.chat_models import ChatZhipuAI
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import PyPDFLoader
import warnings
warnings.filterwarnings('ignore')

# 直接设置API Key
ZHIPU_API_KEY = "2f39319bdd864fc4a41bf6b8eed6efbc.uIsAkRrMwejVTIyc"  # 替换真实密钥

loader = PyPDFLoader('05_机器人工程设计专项赛.pdf')
documents = loader.load()
text_splitter = RecursiveCharacterTextSplitter(chunk_size=1000, chunk_overlap=100)
chunked_documents = text_splitter.split_documents(documents)

vectorstore = Qdrant.from_documents(
    documents=chunked_documents,
    embedding=ZhipuAIEmbeddings(api_key=ZHIPU_API_KEY),
    location=":memory:",
    collection_name="my_documents"
)

retriever = vectorstore.as_retriever()

setup_and_retrieval = RunnableParallel(
    {"context": retriever, "question": RunnablePassthrough()}
)

template = """Answer the question based only on the following context:
{context}

Question: {question}
"""

prompt = ChatPromptTemplate.from_template(template)

model_zhipu = ChatZhipuAI(
    temperature=0.5,
    api_key=ZHIPU_API_KEY,
    model="glm-4"
)

output_parser = StrOutputParser()
chain = setup_and_retrieval | prompt | model_zhipu | output_parser

answer1 = chain.invoke('智能应用专项赛报名时间')
print(answer1)